package com.alibaba.dubbo.remoting.msgcodec;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.math.BigInteger;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.RSAPrivateKeySpec;
import java.security.spec.RSAPublicKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.HashMap;
import java.util.Map;

import javax.crypto.Cipher;


import com.alibaba.dubbo.common.io.Bytes;
import com.alibaba.dubbo.common.io.UnsafeByteArrayOutputStream;

@SuppressWarnings("restriction")
public class RSACodec  {
	public static final String ALGORITHM = "RSA";
	public static final String CIPHER_ALGORITHM = "RSA/ECB/PKCS1Padding";
	public static final String PUBLIC_KEY = "public";
	public static final String PRIVATE_KEY = "private";
	 
	public static Map<String, Object> getKeys() throws NoSuchAlgorithmException,
			InvalidKeySpecException {
		return getKeys(1024);
	}

	@SuppressWarnings("serial")
	public static Map<String, Object> getKeys(int bits)
			throws NoSuchAlgorithmException, InvalidKeySpecException {
		KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance(ALGORITHM);
		keyPairGen.initialize(bits);
		KeyPair keyPair = keyPairGen.generateKeyPair();
		final PrivateKey privateKey = keyPair.getPrivate(); // Extract Private
															// Key
		final PublicKey publicKey = keyPair.getPublic(); // Extract Public Key
		Map<String, Object> result = new HashMap<String, Object>() {
			{
				put(PUBLIC_KEY, publicKey);
				put(PRIVATE_KEY, privateKey);
			}
		};

		return result;
	}

	public static Map<String, String> getKeysAsBase64()
			throws  Exception  {
		return getKeysAsBase64(1024);
	}

	@SuppressWarnings("serial")
	public static Map<String, String> getKeysAsBase64(int bits)
			throws Exception {
		KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance(ALGORITHM);
		keyPairGen.initialize(bits);
		KeyPair keyPair = keyPairGen.generateKeyPair();
		final PrivateKey privateKey = keyPair.getPrivate(); // Extract Private
															// Key
		final PublicKey publicKey = keyPair.getPublic(); // Extract Public Key
		Map<String, String> result = new HashMap<String, String>() {
			{
				put(PUBLIC_KEY, Bytes.bytes2base64(publicKey.getEncoded()));
				put(PRIVATE_KEY,   Bytes.bytes2base64(privateKey.getEncoded()));
			}
		};

		return result;
	}

	/**
	 * 得到公钥
	 * 
	 * @param key
	 *            密钥字符串（经过base64编码）
	 * @throws IOException
	 * @throws NoSuchAlgorithmException
	 * @throws InvalidKeySpecException
	 * @throws Exception
	 */
	static public PublicKey getPublicKeyWithBase64(String key) throws Exception {
		byte[] keyBytes;
		keyBytes =  Bytes.base642bytes(key);

		X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);
		KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
		PublicKey publicKey = keyFactory.generatePublic(keySpec);
		return publicKey;
	}
	
	static public PublicKey getPublicKey(byte []keyBytes) throws Exception {

		X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);
		KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
		PublicKey publicKey = keyFactory.generatePublic(keySpec);
		return publicKey;
	}

	/**
	 * 得到私钥
	 * 
	 * @param key
	 *            密钥字符串（经过base64编码）
	 * @throws Exception
	 */
	public static PrivateKey getPrivateKeyWithBase64(String key)
			throws Exception {
		byte[] keyBytes;
		keyBytes = Bytes.base642bytes(key);

		PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
		KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
		PrivateKey privateKey = keyFactory.generatePrivate(keySpec);
		return privateKey;
	}
	
	public static PrivateKey getPrivateKey(byte[] keyBytes)
			throws Exception {
		PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
		KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
		PrivateKey privateKey = keyFactory.generatePrivate(keySpec);
		return privateKey;
	}

	/**
	 * 使用模和指数生成RSA公钥 注意：【此代码用了默认补位方式，为RSA/None/PKCS1Padding，不同JDK默认的补位方式可能不同，
	 * 如Android默认是RSA/None/NoPadding】
	 * 
	 * @param modulus
	 *            模
	 * @param exponent
	 *            指数
	 * @return
	 */
	public static RSAPublicKey getPublicKey(String modulus, String exponent) {
		try {
			BigInteger b1 = new BigInteger(modulus);
			BigInteger b2 = new BigInteger(exponent);
			KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
			RSAPublicKeySpec keySpec = new RSAPublicKeySpec(b1, b2);
			return (RSAPublicKey) keyFactory.generatePublic(keySpec);
		} catch (Exception e) {
			e.printStackTrace();
			return null;
		}
	}

	/**
	 * 使用模和指数生成RSA私钥 注意：【此代码用了默认补位方式，为RSA/None/PKCS1Padding，不同JDK默认的补位方式可能不同，
	 * 如Android默认是RSA/None/NoPadding】
	 * 
	 * @param modulus
	 *            模
	 * @param exponent
	 *            指数
	 * @return
	 */
	public static RSAPrivateKey getPrivateKey(String modulus, String exponent) {
		try {
			BigInteger b1 = new BigInteger(modulus);
			BigInteger b2 = new BigInteger(exponent);
			KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
			RSAPrivateKeySpec keySpec = new RSAPrivateKeySpec(b1, b2);
			return (RSAPrivateKey) keyFactory.generatePrivate(keySpec);
		} catch (Exception e) {
			e.printStackTrace();
			return null;
		}
	}

	static public byte[] encrypt(byte[] data, RSAPublicKey publicKey) throws Exception{
		Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM);
		cipher.init(Cipher.ENCRYPT_MODE, publicKey);
		int key_len = publicKey.getModulus().bitLength() / 8;
		int max = key_len - 11;
		if(data.length <= max){
			return cipher.doFinal(data);
		}else {
			UnsafeByteArrayOutputStream out = new UnsafeByteArrayOutputStream(4096);
			int offset = 0;
			int len = data.length - offset;
			while(len > 0){
				int l = Math.min(len, max);
				out.write(cipher.doFinal(data, offset, l));
				offset += l;
				len -= l;
			}
			out.flush();
			byte[] result = out.toByteArray();
			out.close();
			return result;
		}
	}
	
	static public int encrypt(InputStream in, OutputStream out, RSAPublicKey publicKey) throws Exception{
		Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM);
		cipher.init(Cipher.ENCRYPT_MODE, publicKey);
		int key_len = publicKey.getModulus().bitLength() / 8;
		int max = key_len - 11;
		byte []data = new byte[max];
		int len;
		int count = 0;
		int total = in.available();
		while((len = in.read(data)) > 0){
			count += len;
			byte []written = cipher.doFinal(data, 0, len) ;
			if(written != null)
			out.write(written);
		}
		out.flush();
		return count;
	}
	
	static public byte[] encrypt(byte[] data, int off, int length, RSAPublicKey publicKey) throws Exception{
		Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM);
		cipher.init(Cipher.ENCRYPT_MODE, publicKey);
		int key_len = publicKey.getModulus().bitLength() / 8;
		int max = key_len - 11;
		if(length <= max){
			return cipher.doFinal(data, off, length);
		}else {
			UnsafeByteArrayOutputStream out = new UnsafeByteArrayOutputStream(4096);
			int offset = off;
			int len = length;
			while(len > 0){
				int l = Math.min(len, max);
				len -= l;
				byte [] written =  cipher.doFinal(data, offset, l);
				if(written != null)
					out.write(written);
				offset += l;
			}
			out.flush();
			byte[] result = out.toByteArray();
			out.close();
			return result;
		}
	}
	
	static public byte[] decrypt(byte[] data, RSAPrivateKey privateKey) throws Exception{
		Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM);
		cipher.init(Cipher.DECRYPT_MODE, privateKey);
		int key_len = privateKey.getModulus().bitLength() / 8;
		int max = key_len;
		if(data.length <= max){
			return cipher.doFinal(data);
		}else {
			UnsafeByteArrayOutputStream out = new UnsafeByteArrayOutputStream(4096);
			int offset = 0;
			int len = data.length - offset;
			while(len > 0){
				int l = Math.min(len, max);
				len -= l;
				byte [] written = cipher.doFinal(data, offset, l);
				if(written != null)
					out.write(written);
				offset += l;
			}
			out.flush();
			byte[] result = out.toByteArray();
			out.close();
			return result;
		}
	}
	
	static public byte[] decrypt(byte[] data, int off, int length, RSAPrivateKey privateKey) throws Exception{
		Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM);
		cipher.init(Cipher.DECRYPT_MODE, privateKey);
		int key_len = privateKey.getModulus().bitLength() / 8;
		int max = key_len;
		if(length <= max){
			return cipher.doFinal(data, off, length);
		}else {
			UnsafeByteArrayOutputStream out = new UnsafeByteArrayOutputStream(4096);
			int offset = off;
			int len = length;
			while(len > 0){
				int l = Math.min(len, max);
				len -= l;
				byte [] written = cipher.update(data, offset, l);
				if(written != null){
					out.write(written);
				}
				out.write(cipher.doFinal(data, offset, l));
				offset += l;
				
			}
			out.flush();
			byte[] result = out.toByteArray();
			out.close();
			return result;
		}
	}
	
	static public int decrypt(InputStream in, OutputStream out, RSAPrivateKey privateKey) throws Exception{
		Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM);
		cipher.init(Cipher.DECRYPT_MODE, privateKey);
		int key_len = privateKey.getModulus().bitLength() / 8;
		int max = key_len;
		byte []data = new byte[max];
		int len;
		int count = 0;
		int total = in.available();
		while ((len = in.read(data)) > 0) {
			count += len;
			byte[] written = cipher.doFinal(data, 0, len);
			if(written != null){
				out.write(written);
			}
		}
		out.flush();
		return count;
	}
	
	static public void main(String []args) throws Exception{
		Map<String, String> keys = getKeysAsBase64();
		System.out.println(keys.get(PUBLIC_KEY));
		System.out.println(keys.get(PRIVATE_KEY));
		String s = "我是中国人abcdegadlkfjalsdfkjalsdkjfalsdkfjalsdkfjalskdfjalsdkfj我是中国人abcdegadlkfjalsdfkjalsdkjfalsdkfjalsdkfjalskdfjalsdkfj我是中国人abcdegadlkfjalsdfkjalsdkjfalsdkfjalsdkfjalskdfjalsdkfj我是中国人abcdegadlkfjalsdfkjalsdkjfalsdkfjalsdkfjalskdfjalsdkfj我是中国人abcdegadlkfjalsdfkjalsdkjfalsdkfjalsdkfjalskdfjalsdkfj";
		byte []data = encrypt(s.getBytes("UTF-8"), (RSAPublicKey) getPublicKeyWithBase64(keys.get(PUBLIC_KEY)) );
		byte []data2 = decrypt(data, (RSAPrivateKey) getPrivateKeyWithBase64(keys.get(PRIVATE_KEY)));
		String d = new String (data2, "UTF-8");
		System.out.println(d);
	}
}
